# 调试/实验脚本
"""
代码尝试专用脚本（空壳模板）
- 用于调试、测试新功能或片段
使用 rich 替换终端输出与进度条。
"""
import time
from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn

from utils.plot_style import print_start, print_end, set_chinese_font, set_plot_style, force_chinese_font

from src.plot import prepare_border_counties_population_data
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.ticker import FuncFormatter
from pathlib import Path
import io
import imageio
import numpy as np
import pandas as pd
import matplotlib.font_manager as fm


def run_tests():
    """在这里放置你的快速试验代码或单元级小片段。"""
    from rich.console import Console
    console = Console()
    console.rule("try.py: 运行示例测试")
    # 非破坏性测试：尝试读取 config 指定的数据文件（若不存在则打印提示）
    try:
        from src.core_utils import load_data
        try:
            df = load_data()
            console.log(f"load_data 成功，rows={len(df)} cols={len(df.columns)}")
            # 仅示例性加载数据并打印行列信息（非破坏性）
            console.log("注意：回归与示例已被移除，请参见 examples/ 目录中的示例脚本。")
        except FileNotFoundError as e:
            console.log(f"load_data 未找到文件：{e}")
        except Exception as e:
            console.log(f"load_data 出错：{e}")
    except Exception:
        console.log("src.core_utils.load_data 不可用：跳过数据加载测试")
    # 演示 rich 进度条
    with Progress(
        SpinnerColumn(style="red"),
        TextColumn("[bold]测试进度"),
        BarColumn(bar_width=None, complete_style="red"),
        TextColumn("{task.completed}/{task.total}"),
    ) as progress:
        task = progress.add_task("test", total=5)
        for _ in range(5):
            time.sleep(0.2)
            progress.advance(task)


def generate_heatmap_from_prepared_data(year: int = 2020, out_dir: str = 'figures', scope: str = 'county') -> str:
    """使用 src.plot.prepare_border_counties_population_data 准备数据并在此处完成绘图与保存。

    返回保存的相对路径字符串。
    """
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    # 准备数据（不在 src.plot 中绘图/保存）
    prepared = prepare_border_counties_population_data(year=year, scope=scope)
    geo = prepared['geo']
    cmap = prepared['cmap']
    vmin = prepared['vmin']
    vmax = prepared['vmax']
    plot_figsize = prepared['figsize']
    plot_dpi = prepared['dpi']
    filename = prepared['filename']
    # filename 由 prepare 返回（已包含中文命名），无需在此追加后缀
    cb_bbox = prepared['colorbar_bbox']
    font_family = prepared['font_family']

    # 应用样式配置（如果 prepare 返回了 cfg_plot，则使用它以保持一致性）
    plot_cfg = prepared.get('cfg_plot') or {}
    set_plot_style(plot_cfg or {})

    # 仅在字体可用时加载中文字体，避免重复或不可用字体导致的警告
    if prepared.get('font_available') and prepared.get('font_family'):
        set_chinese_font(prepared.get('font_family'))
    else:
        # 建议字体不可用或未指定：静默回退到系统默认字体（避免在终端生成多余提示）
        pass

    fig, ax = plt.subplots(figsize=plot_figsize, dpi=plot_dpi)
    geo.plot(column='总人口', cmap=cmap, linewidth=0.4, edgecolor='#444444', ax=ax)
    ax.set_axis_off()

    mappable = mpl.cm.ScalarMappable(norm=mpl.colors.Normalize(vmin=vmin, vmax=vmax), cmap=cmap)
    mappable._A = []
    cax = fig.add_axes(cb_bbox)
    cb = fig.colorbar(mappable, cax=cax, orientation='horizontal')

    def per_ten_thousand(x, pos):
        val = x / 10000.0
        if val >= 10:
            return f"{val:.0f}"
        return f"{val:.1f}"

    cb.ax.xaxis.set_major_formatter(FuncFormatter(per_ten_thousand))
    cb.ax.tick_params(labelsize=9)
    cb.set_label('七普边境县乡镇总人口数（万人）', fontsize=11)

    force_chinese_font(ax, font_family)
    for label in cb.ax.get_xticklabels():
        label.set_fontfamily(font_family)
    cb.ax.xaxis.label.set_fontfamily(font_family)

    fig.subplots_adjust(top=0.92, left=0.06, right=0.98, bottom=0.06)
    out_path = out_dir / filename
    bbox_opt = (prepared['cfg_plot'].get('bbox_inches') if prepared.get('cfg_plot') and prepared['cfg_plot'].get('bbox_inches') else 'tight')
    fig.savefig(out_path, dpi=plot_dpi, bbox_inches=bbox_opt)
    repo_root = Path(__file__).resolve().parents[1]
    try:
        rel_path = out_path.relative_to(repo_root)
    except Exception:
        rel_path = out_path
    from rich.console import Console
    Console().log(f'已保存热力图到 {rel_path}')
    return str(rel_path)


def generate_density_gif(years, out_path: str = 'figures/边境乡镇_mean_density_2000-2020.gif', column: str = 'mean', cmap: str = 'Reds', dpi: int = 150, figsize=(10, 12), scope: str = 'county', duration: float = 0.6, population_parquet: str = 'data/边境城市乡镇人口密度面板数据.parquet'):
    """为给定年份列表渲染每年人口密度字段（默认为 'mean'）并合成 GIF。

    说明：使用 prepare_border_counties_population_data 来准备每年的 GeoDataFrame，
    但不保存中间静态图像，直接把 matplotlib 图转为 RGBA 数组并传给 imageio。
    """
    out_dir = Path('figures')
    out_dir.mkdir(parents=True, exist_ok=True)
    frames = []

    # 为了保证色标一致性，先计算全时间段针对所选列的 vmin/vmax
    vmin = None
    vmax = None
    geos = {}
    for y in years:
        prep = prepare_border_counties_population_data(year=y, scope=scope, population_parquet=population_parquet)
        if not prep:
            continue
        geos[y] = prep
        geo = prep['geo']
        # 以所选列为准（回退到 'mean' 或 '总人口'）
        if column in geo.columns:
            ser = pd.to_numeric(geo[column], errors='coerce').fillna(0)
        elif 'mean' in geo.columns:
            ser = pd.to_numeric(geo['mean'], errors='coerce').fillna(0)
        else:
            ser = pd.Series(dtype=float)
        if not ser.empty:
            yr_min = float(ser.min())
            yr_max = float(ser.max())
            if vmin is None or yr_min < vmin:
                vmin = yr_min
            if vmax is None or yr_max > vmax:
                vmax = yr_max

    if not geos:
        from rich.console import Console
        Console().log('未找到任何年份的数据，GIF 生成终止。')
        return None

    # 如果全程恒定（vmin == vmax），扩大一个微小范围以避免颜色映射失败
    if vmin is None or vmax is None:
        vmin = 0.0
        vmax = 1.0
    if vmin == vmax:
        vmax = vmin + 1e-6

    # 尝试找到一个可用的中文字体
    preferred_cn = ['Noto Sans CJK SC', 'SimHei', 'PingFang', 'Heiti SC', 'WenQuanYi Zen Hei', 'Microsoft YaHei', 'STHeiti']
    available_font_name = None
    try:
        names = {f.name for f in fm.fontManager.ttflist}
        for cand in preferred_cn:
            if cand in names:
                available_font_name = cand
                break
    except Exception:
        available_font_name = None

    # 渲染每一年为一帧
    for y in sorted(geos.keys()):
        prep = geos[y]
        geo = prep['geo']
        plot_figsize = prep['figsize']
        plot_dpi = prep['dpi']
        font_family = prep.get('font_family')

        # 样式
        set_plot_style(prep.get('cfg_plot') or {})
        if prep.get('font_available') and font_family:
            set_chinese_font(font_family)

        fig, ax = plt.subplots(figsize=figsize or plot_figsize, dpi=plot_dpi)
        # 绘制指定列作为热力图
        if column not in geo.columns:
            # 如果没有该列，尝试使用 'mean' 或 '总人口' 作为回退
            col = 'mean' if 'mean' in geo.columns else '总人口'
        else:
            col = column

        geo.plot(column=col, cmap=cmap, linewidth=0.2, edgecolor='#444444', ax=ax, vmin=vmin, vmax=vmax)
        ax.set_axis_off()

        mappable = mpl.cm.ScalarMappable(norm=mpl.colors.Normalize(vmin=vmin, vmax=vmax), cmap=cmap)
        mappable._A = []
        cax = fig.add_axes(prep['colorbar_bbox'])
        cb = fig.colorbar(mappable, cax=cax, orientation='horizontal')

        def per_ten_thousand(x, pos):
            val = x / 10000.0
            if val >= 10:
                return f"{val:.0f}"
            return f"{val:.1f}"

        try:
            cb.ax.xaxis.set_major_formatter(FuncFormatter(per_ten_thousand))
            cb.ax.tick_params(labelsize=9)
            cb.set_label('人口密度（单位根据数据）', fontsize=11)
        except Exception:
            pass

        force_chinese_font(ax, font_family)

        # 在内存中保存为 RGBA 图像并加入 frames
        buf = io.BytesIO()
        fig.savefig(buf, format='png', dpi=plot_dpi, bbox_inches='tight')
        plt.close(fig)
        buf.seek(0)
        img = imageio.v2.imread(buf)
        frames.append(img)

    # 写入 GIF
    gif_path = Path(out_path)
    imageio.mimsave(gif_path, frames, format='GIF', duration=duration)
    from rich.console import Console
    try:
        repo_root = Path(__file__).resolve().parents[1]
        rel = gif_path.relative_to(repo_root)
    except Exception:
        rel = gif_path
    Console().log(f'已保存 GIF 到 {rel} (frames={len(frames)})')
    return str(rel)


if __name__ == "__main__":
    start_time = print_start('生成热力图（try.py）')
    # 直接生成热力图并保存
    try:
        # 生成边境县范围热力图（原先行为）
        generate_heatmap_from_prepared_data(year=2020, out_dir='figures', scope='county')
        # 生成边境城市范围热力图（新的需求）
        generate_heatmap_from_prepared_data(year=2020, out_dir='figures', scope='city')
    except Exception as e:
        from rich.console import Console
        Console().log(f'生成热力图出错：{e}')
    print_end(start_time, '生成热力图（try.py）')

